package me.shan.sort.merge;

/**
 * <a href="https://leetcode.cn/problems/count-of-range-sum/">区间和的个数</a>
 */
public class CountOfRangeSum {

    public static int countRangeSum(int[] nums, int lower, int upper) {
        if (nums == null || nums.length == 0) {
            return 0;
        }
        // 求前缀和数组
        long[] sums = new long[nums.length];
        sums[0] = nums[0];
        for (int i = 1; i < nums.length; i++) {
            sums[i] = sums[i - 1] + nums[i];
        }
        // 直接调用归并获取结果
        return mergeCount(sums, 0, sums.length - 1, lower, upper);
    }

    private static int mergeCount(long[] sums, int left, int right, int lower, int upper) {
        // base case
        if (left == right) {
            if (sums[left] >= lower && sums[left] <= upper) {
                return 1;
            } else {
                return 0;
            }
        }

        int mid = left + (right - left) / 2;
        int leftAns = mergeCount(sums, left, mid, lower, upper);
        int rightAns = mergeCount(sums, mid + 1, right, lower, upper);
        int mergeAns = merge(sums, left, mid, right, lower, upper);
        return leftAns + rightAns + mergeAns;
    }

    private static int merge(long[] sums, int left, int mid, int right, int lower, int upper) {
        int ans = 0;
        // 求解
        int windowL = left;
        int windowR = left;
        int j = mid + 1;
        while (j <= right) {
            long ansLower = sums[j] - upper;
            long ansUpper = sums[j] - lower;
            while (sums[windowL] < ansLower && windowL <= mid) {
                windowL++;
            }
            while (sums[windowR] <= ansUpper && windowR <= mid) {
                windowR++;
            }
            ans += windowR - windowL;
            j++;
        }
        // 归并
        long[] temp = new long[right - left + 1];
        int a = left, b = mid + 1, c = 0;
        while (a <= mid && b <= right) {
            if (sums[a] <= sums[b]) {
                temp[c++] = sums[a++];
            } else {
                temp[c++] = sums[b++];
            }
        }
        while (a <= mid) {
            temp[c++] = sums[a++];
        }
        while (b <= right) {
            temp[c++] = sums[b++];
        }
        System.arraycopy(temp, 0, sums, left, temp.length);
        return ans;
    }

    /**
     * 暴力递归发
     * @param numbs 原始数组
     * @param lower 下限
     * @param upper 上限
     * @return 答案
     */
    public static int countRangeSum2(int[] numbs, int lower, int upper) {
        if (numbs == null || numbs.length == 0) {
            return 0;
        }
        long[] sums = new long[numbs.length];
        sums[0] = numbs[0];
        for (int i = 1; i < numbs.length; i++) {
            sums[i] = sums[i - 1] + numbs[i];
        }
        int ans = 0;
        for (int len = 1; len <= numbs.length; len++) {
            for (int i = 0; i <= numbs.length - len; i++) {
                long num = sum(sums, i, len);
                if (num >= lower && num <= upper) {
                    ans++;
                }
            }
        }
        return ans;
    }

    private static long sum(long[] sums, int index, int len) {
        if (index == 0) {
            return sums[len - 1];
        }
        return sums[index + len - 1] - sums[index - 1];
    }

    public static void main(String[] args) {
        int[] numbs = {-3, 1, 2, -2, 2, -1};
        System.out.println(countRangeSum(numbs, -3, -1));
        System.out.println(countRangeSum2(numbs, -3, -1));
    }

}
